#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
修复项目ID，确保符合Milvus集合命名规范
"""

import os
import asyncio
import sys
import re

# 添加项目根目录到路径
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.future import select
from sqlalchemy import update

from app.core.logging import setup_logging, get_logger
from app.core.utils import validate_collection_name
from app.models.database import Project
from app.repositories.mysql import get_db_session, init_db
from app.repositories.milvus import milvus_client
from app.repositories.vector import VectorRepository

# 创建日志记录器
logger = get_logger("fix_project_ids")


async def fix_project_ids():
    """修复项目ID，确保符合Milvus集合命名规范"""
    try:
        # 连接到Milvus
        await milvus_client.connect()
        
        # 获取所有项目
        projects = []
        async for session in get_db_session():
            result = await session.execute(select(Project))
            projects = result.scalars().all()
            
            # 检查并修复项目ID
            for project in projects:
                project_id = project.id
                logger.info(f"检查项目ID: {project_id}")
                
                # 检查项目ID是否符合Milvus集合命名规范
                if not validate_collection_name(project_id):
                    # 生成新的项目ID
                    new_project_id = project_id.replace("-", "_")
                    logger.info(f"项目ID不符合规范，将修改为: {new_project_id}")
                    
                    # 更新项目ID
                    await session.execute(
                        update(Project)
                        .where(Project.id == project_id)
                        .values(id=new_project_id)
                    )
                    
                    # 提交更改
                    await session.commit()
                    logger.info(f"已更新项目ID: {project_id} -> {new_project_id}")
                else:
                    logger.info(f"项目ID符合规范，无需修改: {project_id}")
            
            break  # 只需要执行一次查询
        
        logger.info("项目ID修复完成")
    except Exception as e:
        logger.error(f"修复项目ID失败: {str(e)}")
        raise


async def main():
    """主函数"""
    # 设置日志
    setup_logging()
    
    # 初始化数据库
    await init_db()
    
    # 修复项目ID
    await fix_project_ids()


if __name__ == "__main__":
    # 运行异步主函数
    asyncio.run(main())